import pandas as pd
import numpy as np
import matplotlib

# 使用非交互式后端，避免GUI相关错误
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import seaborn as sns
from matplotlib.colors import LinearSegmentedColormap
import matplotlib.font_manager as fm

# 设置中文字体支持
plt.rcParams['font.sans-serif'] = ['SimHei']  # 用来正常显示中文标签
plt.rcParams['axes.unicode_minus'] = False  # 用来正常显示负号

# 读取Excel文件
file_path = 'frog_similarity_results.xlsx'
df = pd.read_excel(file_path, sheet_name='Sheet1')

# 获取唯一的图片名称并排序
image_names = sorted(set(df['一号实验对象']).union(set(df['二号实验对象'])),
                     key=lambda x: int(x.split('.')[0]))

# 创建映射字典
image_to_index = {name: idx for idx, name in enumerate(image_names)}
n = len(image_names)

# 创建自定义颜色映射
colors = ["#1E90FF", "#FFFFFF", "#FF6347"]  # 蓝 -> 白 -> 红
cmap = LinearSegmentedColormap.from_list("custom_cmap", colors, N=256)

# 创建四个独立的图形，每个指标一个
metrics = [
    ('ResNet50 Similarity', 'ResNet50相似度矩阵'),
    ('VGG16 Similarity', 'VGG16相似度矩阵'),
    ('SIFT Similarity', 'SIFT相似度矩阵'),
    ('相似度（去掉MSE，预训练）', '综合相似度矩阵')
]

# 为每个指标创建矩阵并绘制
for metric, title in metrics:
    # 创建空矩阵
    matrix = np.zeros((n, n))

    # 填充矩阵
    for _, row in df.iterrows():
        row_idx = image_to_index[row['一号实验对象']]
        col_idx = image_to_index[row['二号实验对象']]
        value = row[metric]
        matrix[row_idx, col_idx] = value
        matrix[col_idx, row_idx] = value  # 确保对称

    # 对角线设为1或100（根据指标类型）
    np.fill_diagonal(matrix, 100 if metric == '相似度（去掉MSE，预训练）' else 1)

    # 创建DataFrame用于绘图
    matrix_df = pd.DataFrame(matrix, index=image_names, columns=image_names)

    # 创建图形
    plt.figure(figsize=(12, 10))

    # 设置vmax参数
    vmax = 100 if metric == '相似度（去掉MSE，预训练）' else 1

    # 绘制热力图
    ax = sns.heatmap(
        matrix_df,
        annot=True,
        fmt=".2f" if metric != '相似度（去掉MSE，预训练）' else ".1f",
        cmap=cmap,
        center=vmax / 2,
        vmin=0,
        vmax=vmax,
        square=True,
        cbar_kws={'label': metric}
    )

    # 设置标题和标签
    plt.title(f'{title}', fontsize=16, pad=20)
    plt.xlabel('图片名称', fontsize=12)
    plt.ylabel('图片名称', fontsize=12)

    # 调整x轴标签
    plt.xticks(rotation=45, ha='right', fontsize=8)
    plt.yticks(fontsize=8)

    # 调整布局
    plt.tight_layout()

    # 保存图像
    filename = f'frog_similarity_{metric.split()[0]}.png'
    plt.savefig(filename, dpi=300, bbox_inches='tight')
    print(f"已保存: {filename}")

    # 关闭图形，释放内存
    plt.close()

print("所有矩阵图已生成完成！")